import json
import os


# {
#     "id": 5,
#     "query": "什么是裸金属",
#     "document": "director",
#     "related_file_path": [
#         "data/director/配置指南/23-发布裸金属.txt"
#     ]
# }


class RagItem:

    def __init__(self,id:int,query:str,file_path_list:list[str]):
        self.id = id
        self.query = query
        self.file_path_list = file_path_list
    
    def __repr__(self) -> str:
        content = """
ID: {}
Query: {}
File Path List: 
    {}
""".format(
    self.id,
    self.query,
    "\n    ".join(self.file_path_list)
)
        return content


class AnswerLoader:

    # 递归检索的结果/最终结果00/retriv-results.json

    def __init__(self,file_path:str):
        self.file_path = file_path
        self.data = []
        with open(self.file_path,"r",encoding="utf-8") as file:
            content = file.read()
            self.data = json.loads(content)
        # dict
        self.query_dict:dict[int,RagItem] = dict()
        self._load()
        pass

    def _add_ragitem(self,id:int,query:str,file_path_list:list[str]):
        """
        添加RagItem映射
        """
        if self.query_dict.__contains__(id):
            raise Exception(f"ID:{id}已经存在")
        rag = RagItem(id,query,file_path_list)
        self.query_dict[id] = rag
        # print(f"{rag}")

    def _load(self):
        for item in self.data:
            id = item["id"]
            query = item["query"]
            nodes = item["nodes"]
            path_list = []
            for node in nodes:
                path_list.append(node["metadata"]["file_path"])
            path_list = list(set(path_list)) # 去重复
            # 添加到rag
            self._add_ragitem(id,query,path_list)
        pass

    def IDs(self) -> list[int]:
        return list(self.query_dict.keys())

    def get_ragitem(self,id:int):
        item = self.query_dict.get(id,None)
        if item is None:
            return None
        return item



class GtLoader:

    def __init__(self,dir_path:str):
        self.reset(dir_path)
    
    def reset(self,dir_path:str):
        self.query_dict:dict[int,RagItem] = dict()
        self._loadgt(dir_path)
    
    def _add_ragitem(self,id:int,query:str,file_path_list:list[str]):
        """
        添加RagItem映射
        """
        if self.query_dict.__contains__(id):
            raise Exception(f"ID:{id}已经存在")
        rag = RagItem(id,query,file_path_list)
        self.query_dict[id] = rag
        # print(f"{rag}")

    def get_ragitem(self,id:int):
        item = self.query_dict.get(id,None)
        if item is None:
            return None
        return item
    
    def IDs(self) -> list[int]:
        return list(self.query_dict.keys())

    def _loadgt(self, dir_path:str):
        # 收集所有的json文件
        all = []
        for file in os.listdir(dir_path):
            if file.endswith(".json"):
                all.append(file)
        # 读取文件
        for filename in all:
            try:
                filename = str(filename)
                file_id = int(filename.split(".json")[0])
            except Exception as e:
                print(f"文件名不符合规范: {filename}")
                continue
            file_path = os.path.join(dir_path,filename)
            with open(file_path,"r",encoding="utf-8") as file:
                data = json.loads(file.read())
                id = data["id"]
                query = data["query"]
                related_path = data["related_file_path"]
                self._add_ragitem(id,query,related_path)

